import numpy as np
from astropy.time import Time
from astropy import constants
import astropy.units as u
from astropy.coordinates import Longitude, Latitude, Angle
from astropy import coordinates
import utils

def newton_generation(initial_value, e_1, M_1, tolerance):
    """
    牛顿迭代求解开普勒方程 E - e sin E = M
    :param initial_value: 开始迭代的值，可以为任意值，下面代码中取了上一个时刻的E为初始值
    :param e_1: 加了_1后缀是为了和下面代码中的e和M进行区分
    :param M_1: 同上
    :param tolerance: 迭代到误差小于什么时候停止，下面代码中取了1e-6
    :return: 计算结果
    """
    difference = 2 * tolerance
    old_value = initial_value
    new_value = initial_value
    while difference > tolerance:
        new_value = old_value - (old_value - e_1 * np.sin(old_value) * u.rad - M_1) / (1 - e_1 * np.cos(old_value))
        difference = abs(old_value - new_value)
        old_value = new_value

    return new_value

# 测站的坐标，参考了作业3中的样例代码
site_lon = 118.78 * u.deg
site_lat = 32.07 * u.deg
site_height = 0 * u.m
site_coord = coordinates.EarthLocation.from_geodetic(site_lon, site_lat, 0 * u.m).to_value('m')
# 转为array，方便后续计算
site_coord_np = np.array([site_coord[0], site_coord[1], site_coord[2]]) * u.m


utc_time = '2022-03-16T11:06:00'

# 下面计算轨道根数，先得到目标的初始时刻位置和速度
object_initial_location = np.array([-915731, 5958096, 3050295]) * u.m
object_initial_speed = np.array([-6090.2, -2849.3, 3718.7]) * u.m / u.s

# 由目标初始位置和速度得到h
h_vector = np.outer(object_initial_location, object_initial_speed)
h_scalar = np.linalg.norm(h_vector)

# 由h矢量得到方向向量R
R_vector = h_vector / h_scalar

# 再由方向向量得到i和Ω
i = np.arccos(R_vector[-1])
omega = np.arcsin(R_vector[0] / np.sin(i))

# 引入μ，参考书本附录，模型是WGS84, 因为测站的坐标转换中，第32行代码里from_geodetic，默认选择了WGS84模型
mu = 398600.4418 * 1e9 * (u.m ** 3) / ((u.s) ** 2)

# 计算目标初始时刻位置和速度的大小
r_scalar = np.linalg.norm(object_initial_location)
v_scalar = np.linalg.norm(object_initial_speed)

# 计算轨道半长轴
a = 1 / (2 / r_scalar - v_scalar ** 2 / mu)

# 计算E, e, M
E = np.arctan2(np.dot(object_initial_speed, object_initial_location) / np.sqrt(mu * a), (1 - r_scalar / a))
e = (1 - r_scalar / a) / np.cos(E)
M = E - e * np.sin(E) * u.rad

# 计算向量P和Q，用于后续得到目标位置和速度的矢量
P_vector = np.cos(E) / r_scalar * object_initial_location - np.sqrt(a / mu) * np.sin(E) * object_initial_speed
Q_vector = (np.sin(E) / r_scalar * object_initial_location + np.sqrt(a / mu) * (np.cos(E) - e) * object_initial_speed)

# 计算ω， 但后续没用到
w = np.arctan2(P_vector[-1], Q_vector[-1])

# 计算n，用于更新M，因为dM = n dt
n = np.sqrt(mu / a ** 3)

# flag 用于记录满足了多少个可见条件，一共3个
flag = 0

# available 用于记录当前是否已经可见
available = 0

# 问题一的输出极移差别的一个参数
print_difference_times = 0

# print(np.pi * 2 / n)
# 假定一个小时时间，轨道的周期可由2π / n 得到， 计算出是5531秒
for i in range(3600):
    time = Time(utc_time, scale='utc') + i * u.s

    # 更新轨道根数M和E
    M = M + n * u.s * u.rad
    E = newton_generation(E, e, M, 1e-6 * u.rad)

    # 计算目标的位置向量和速度向量
    r_vector = a * (np.cos(E) - e) * P_vector + a * (np.sqrt(1 - e ** 2) * np.sin(E) * Q_vector)
    r_scalar = np.linalg.norm(r_vector)
    speed_vector = n * a ** 2 / r_scalar * (-np.sin(E) * P_vector + np.sqrt(1 - e ** 2) * np.cos(E) * Q_vector)

    # 用作业2的代码进行参考系的转换，极移量级很小，无需考虑
    transform_matrix = utils.CoordinatesTransformEquinoxBase(time)
    r_vector_earth_location = transform_matrix.earth_rotation_matrix() @ r_vector
    speed_vector_earth_location = transform_matrix.earth_rotation_matrix() @ speed_vector

    # r_vector_earth_location_p = transform_matrix.polar_motion_matrix() @ transform_matrix.earth_rotation_matrix() @ r_vector
    # speed_vector_earth_location_p = transform_matrix.polar_motion_matrix() @ transform_matrix.earth_rotation_matrix() @ speed_vector
    #
    # if print_difference_times == 0:
    #     print_difference_times += 1
    #     print(r_vector_earth_location - r_vector_earth_location_p)
    #     print(speed_vector_earth_location - speed_vector_earth_location_p)


    # 条件1，卫星在测站视野范围内：h > h_0
    h_0 = 5.0 * u.deg
    theta = np.arccos(np.dot(site_coord_np, r_vector_earth_location) / (np.linalg.norm(site_coord_np) * np.linalg.norm(r_vector_earth_location)))
    theta_0 = np.arccos(np.linalg.norm(site_coord_np) / np.linalg.norm(r_vector_earth_location) * np.cos(h_0)) - h_0
    if theta < theta_0:
        flag += 1


    # 条件2，夜晚，太阳落山
    beta_0 = 6.0 * u.deg
    # 计算太阳坐标，此时为GCRS参考系下的坐标，需要转到ITRS坐标
    sun_coord = coordinates.get_sun(time).transform_to(coordinates.ITRS)
    sun_x = sun_coord.x.to_value('m')
    sun_y = sun_coord.y.to_value('m')
    sun_z = sun_coord.z.to_value('m')
    # 转为array，方便后续计算
    sun_location = np.array([sun_x, sun_y, sun_z]) * u.m

    beta = np.arccos(np.dot(site_coord_np, sun_location) / (np.linalg.norm(site_coord_np) * np.linalg.norm(sun_location)))
    if beta > 90.0 * u.deg + beta_0:
        flag += 1


    # 条件3，在地影外
    psi = np.arccos(np.dot(r_vector_earth_location, sun_location) / (np.linalg.norm(r_scalar) * np.linalg.norm(sun_location)))
    r_edge_scalar = constants.R_earth / np.sin(psi)
    # print(constants.R_earth)
    if r_scalar > r_edge_scalar:
        flag += 1

    # 如果flag == 3 说明三个可见性条件都满足，就输出时间，并把记录当前可见性的available置为1
    if flag == 3:
        if available == 0:
           print(time)
        available = 1

    # 如果上一秒可见，但现在不满足3个可见性条件，就输出现在的时间，并结束循环
    if available == 1 and flag < 3:
        print(time)
        break
    flag = 0

aaa = 1